import torch

from ljp.dataset import medmnist
import torchvision


def get_transforms(mean=[.5], std=[.5]):
    return torchvision.transforms.transforms.Compose([
        torchvision.transforms.transforms.ToTensor(),
        torchvision.transforms.transforms.Normalize(mean=mean, std=std)
    ])


class MEDMNIST:

    def __init__(self, flag='pathmnist', image_size=28, batch_size=128, dataroot='/datahub/medmnist/', mean=[.5],
                 std=[.5]):
        self.info = medmnist.INFO[flag]
        self.dataroot = dataroot
        self.batch_size = batch_size
        self.inchannels = self.info['n_channels']
        self.num_classes = len(self.info['label'])
        self.img_sz = 28
        self.mean = mean
        self.std = std

        DataClass = getattr(medmnist, self.info['python_class'])
        train = DataClass(split="train", root=self.dataroot, transform=get_transforms(mean=mean, std=std))
        # print(train)
        val = DataClass(split="val", root=self.dataroot, transform=get_transforms(mean=mean, std=std))
        test = DataClass(split="test", root=self.dataroot, transform=get_transforms(mean=mean, std=std))
        self.trainloader = torch.utils.data.DataLoader(dataset=train, batch_size=self.batch_size, shuffle=True)
        self.valloader = torch.utils.data.DataLoader(dataset=val, batch_size=256, shuffle=False)
        self.testloader = torch.utils.data.DataLoader(dataset=test, batch_size=256, shuffle=False)
        self.desc = f'{self.info["python_class"]}[bz{self.batch_size}iz{self.img_sz}]'


class MEDMNISTELM:

    def __init__(self, flag='pathmnist', dataroot='/data/medmnist/', mean=[.5], std=[.5]):
        self.info = medmnist.INFO[flag]
        self.dataroot = dataroot
        self.inchannel = self.info['n_channels']
        self.num_classes = len(self.info['label'])
        self.img_sz = 28
        self.mean = mean
        self.std = std

        DataClass = getattr(medmnist, self.info['python_class'])
        train = DataClass(split="train", root=self.dataroot, transform=get_transforms(mean=mean, std=std))
        # print(train)
        val = DataClass(split="val", root=self.dataroot, transform=get_transforms(mean=mean, std=std))
        test = DataClass(split="test", root=self.dataroot, transform=get_transforms(mean=mean, std=std))
        self.trainloader = torch.utils.data.DataLoader(dataset=train, batch_size=self.info['n_samples']['train'],
                                                       shuffle=True)
        self.valloader = torch.utils.data.DataLoader(dataset=val, batch_size=self.info['n_samples']['val'],
                                                     shuffle=False)
        self.testloader = torch.utils.data.DataLoader(dataset=test, batch_size=self.info['n_samples']['test'],
                                                      shuffle=False)
        self.desc = f'{self.info["python_class"]}'


if __name__ == '__main__':
    a = MEDMNIST('pathmnist', batch_size=128)
    for inputs, targets in a.valloader:
        print(targets.shape)
        print(len(targets.shape))
